import functools
import http
import json

from starlette.middleware.base import RequestResponseEndpoint, BaseHTTPMiddleware, DispatchFunction
from starlette.requests import Request
from starlette.responses import Response
import typing
from typing import Optional
from pydantic import BaseModel, Field
from starlette.datastructures import Headers
from starlette.types import ASGIApp, Message, Receive, Scope, Send


class BaseMiddlewareNoResponse:
    '''
    # 这个中间件的话，如果需要在内部-0---读取请求体内容，需要再最终添加，也就是需要在最前面的执行
    # 也就是需要在所有的中间件的最后面再去注册，不能放在其他的前面
    # 如果需要在内部消费   body = await self.request.body()，则需要开启is_proxy=True
    '''
    pass

    def __init__(self, app: ASGIApp, is_proxy=True) -> None:
        self.app = app
        self.is_proxy = is_proxy
        self.request: typing.Optional[Request] = None

    def bind_to_request_state(self, request: Request, **kwargs):
        """
        Takes in a set of kwargs and binds them to gziprequest state
        """
        for key, value in kwargs.items():
            setattr(request.state, key, value)

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] != "http":
            await self.app(scope, receive, send)
            return
        # 解决读取BODY问题
        if self.is_proxy:
            receive_ = await receive()

            async def receive():
                return receive_
        # 解析当前的请求体
        self.request = Request(scope, receive=receive)
        # 自动传参，如果对于send有需要重写的需求，则需要进行重写
        send = functools.partial(self.send, send=send, request=self.request)
        # 自定义回调函数，可以自己进行重写实现具体的业务逻辑
        response = await self.before_request(self.request) or self.app
        await response(self.request.scope, receive, send)
        await self.after_request(self.request)

    async def send(self, message: Message, send: Send, request: Request) -> None:
        """重写send方法【不重写则默认使用原来的】"""
        return await send(message)

    async def before_request(self, request: Request) -> [Response, None]:
        """如果需要修改请求信息，可直接重写此方法"""
        return self.app

    async def after_request(self, request: Request) -> [Response, None]:
        """请求后的处理【记录请求耗时等，注意这里没办法对响应结果进行处理】"""
        return None

    async def get_body(self):
        """获取请求BODY，实现使用代理方式解析读，解决在中间件中火球Body的问题"""
        body = await  self.request.body()
        return body

    async def get_json(self):
        """获取json请求参数"""
        return json.loads(await self.get_body())


class BaseMiddlewareHasResponse:
    '''
    此类的中间件可以接续读取返回的响应报文
    使用上：
    core_app.core_middleware('http')(BaseMiddlewareHasResponse())
    '''
    pass

    def __init__(self, app: ASGIApp) -> None:
        self.app = app
        self.request: typing.Optional[Request] = None

    async def before_request(self, request: Request) -> [Response, None]:
        """如果需要修改请求信息，可直接重写此方法"""
        pass

    async def after_request(self, request: Request, res: Response = None) -> [Response, None]:
        """请求后的处理【记录请求耗时等，注意这里没办法对响应结果进行处理】"""
        return res

    async def get_body(self):
        """获取请求BODY，实现使用代理方式解析读，解决在中间件中火球Body的问题"""
        body = await  self.request.body()
        return body

    async def get_json(self):
        """获取json请求参数"""
        return json.loads(await self.get_body())

    async def __call__(
            self,
            request: Request,
            call_next: RequestResponseEndpoint,
            *args,
            **kwargs
    ):
        try:
            # 下一个响应报文内容
            # 解析当前的请求体
            self.request = request
            await self.before_request(self.request) or self.app
            response = await call_next(request)
        except Exception as ex:
            # 解析响应报文的body异常信息
            response_body = bytes(http.HTTPStatus.INTERNAL_SERVER_ERROR.phrase.encode())
            # 生成异常报文内容
            response = Response(content=response_body, status_code=http.HTTPStatus.INTERNAL_SERVER_ERROR.real,
                                )
        else:
            response_body = b''
            # 解析读取对应的响应报文内容型芯
            async for chunk in response.body_iterator:
                response_body += chunk
            # 响应体最终响应报文内容
            response = Response(content=response_body,
                                status_code=response.status_code,
                                headers=dict(response.headers),
                                media_type=response.media_type
                                )
            await self.after_request(self.request, response)
        return response



class BaseHandlerMiddleware(BaseHTTPMiddleware):

    """Middleware to dispatch modified response"""

    def __init__(self, app: ASGIApp,
                 dispatch: DispatchFunction = None, handler: typing.Callable = None) -> None:
        super().__init__(app, dispatch=dispatch)
        self.handler = handler

    async def dispatch(
            self,request: Request,call_next: RequestResponseEndpoint) -> Response:
        request.state.background = None
        response = await call_next(request)
        # 执行响应报文中的后台任务
        if request.state.background:
            # 处理所有的额相关的fastapi BackgroundTasks的任务，它是在响应报文之中最终被执行处理的
            response.background = request.state.background
        # 处理所有的额相关的fastapi BackgroundTasks的任务，它是在响应报文之中最终被执行处理的
        # 如响应报文的中源码的调用：
        # def __init__(
        #         self,
        #         content: typing.Any = None,
        #         status_code: int = 200,
        #         headers: typing.Optional[typing.Mapping[str, str]] = None,
        #         media_type: typing.Optional[str] = None,
        #         background: typing.Optional[BackgroundTask] = None,
        # ===========================================
        # ) -> None:
        #     async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        #         await send(
        #             {
        #                 "type": "http.response.start",
        #                 "status": self.status_code,
        #                 "headers": self.raw_headers,
        #             }
        #         )
        #         await send({"type": "http.response.body", "body": self.body})
        #         if self.background is not None:
        #             # for iasa in BackgroundTask
        #             await self.background()
        return response

# ======================================
class ResponseInfo(BaseModel):
    headers: Optional[Headers] = Field(default=None, title="Response header")
    body: str = Field(default="", title="Response body")
    status_code: Optional[int] = Field(default=None, title="Status code")

    class Config:
        arbitrary_types_allowed = True


class BaseResponseMiddleware:
    '''
    # 这个中间件的话，如果需要在内部-0---读取请求体内容，需要再最终添加，也就是需要在最前面的执行
    # 也就是需要在所有的中间件的最后面再去注册，不能放在其他的前面,
    # 如果需要在内部消费   body = await self.request.body()，则需要开启is_proxy=True
    '''

    def __init__(self, app: ASGIApp, is_proxy=True) -> None:
        self.app = app
        self.is_proxy = is_proxy
        self.request: typing.Optional[Request] = None

    async def get_body(self):
        """获取请求BODY，实现使用代理方式解析读，解决在中间件中火球Body的问题"""
        body = await self.request.body()
        return body

    async def get_json(self):
        """获取json请求参数"""
        body = await self.get_body()
        return json.loads(body) if body else None

    async def before_request(self, request: Request) -> [Response, None]:
        """如果需要修改请求信息，可直接重写此方法"""
        pass

    async def after_request(self, request: Request, res: Response = None) -> [Response, None]:
        """请求后的处理【记录请求耗时等，注意这里没办法对响应结果进行处理】"""
        pass

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] != "http":
            return await self.app(scope, receive, send)

        # 解决读取BODY问题
        if self.is_proxy:
            receive_ = await receive()

            async def receive():
                return receive_

        # 这里考虑直接读取一次body然后保存到对应的上下文中

        # 解析当前的请求体
        self.request = Request(scope, receive=receive)
        # 解析报文体内容
        response_info = ResponseInfo()
        # 自定义回调函数，可以自己进行重写实现具体的业务逻辑
        await self.before_request(self.request) or self.app
        # 下一个循环体
        # 从响应报文的处理机制可以启发得到下面的处理信息
        #  async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        #         await send(
        #             {
        #                 "type": "http.response.start",
        #                 "status": self.status_code,
        #                 "headers": self.raw_headers,
        #             }
        #         )
        #         await send({"type": "http.response.body", "body": self.body})
        #         if self.background is not None:
        #             # print("当前的后台任务是不是",self.background)
        #             await self.background()

        async def _next_send(message: Message) -> None:
            if message.get("type") == "http.response.start":
                response_info.headers = Headers(raw=message.get("headers"))
                response_info.status_code = message.get("status")
            # 解析响应体内容信息
            elif message.get("type") == "http.response.body":
                if body := message.get("body"):
                    # response_info.body += body.decode("utf8")
                    response_info.body += body.decode("utf-8", errors="ignore")
                # print('响应头吧', response_info)
                response = Response(content=response_info.body,
                                    status_code=response_info.status_code,
                                    headers=dict(response_info.headers)
                                    )
                await self.after_request(self.request, response)
            await send(message)

        await self.app(scope, receive, _next_send)
